from ultralytics import YOLO

if __name__ == '__main__':
    # 第二阶段配置 微调
    stage1_path = f"runs/detect/train18/weights/best.pt"
    stage2_config = {
        "data": "ultralytics/cfg/datasets/helmet4.yaml",
        "epochs": 50,
        "imgsz": 800,
        "batch": 16,
        "lr0": 0.00005,
        "optimizer": "AdamW",
        # "resume": True,
        "project": "runs/detect",
        "device": 0,
        "workers": 0,
        # "patience": 5,  # 微调阶段早停耐心值减小
        "erasing": 0.4,  # 保持与第一阶段一致
    }

    # 执行第二阶段微调
    model = YOLO(stage1_path)
    model.train(**stage2_config)